import json

from fastapi import FastAPI, Request
from gunicorn.app.base import BaseApplication
from pydantic import BaseModel

from check_npu_status import check_multi_npu_not_used

app = FastAPI()

MACHINE_POOL_FILE="machine_pool.json"
STATUS = dict()

class AskNodesRequest(BaseModel):
    dynamic_nodes: bool = True # 动态请求节点
    ask_hosts: list # dynamic_nodes为False时生效，查询目标节点是否空闲
    node_number: int # dynamic_nodes为False时生效，查询空闲的目标个数节点
    pool_id: str


class MachineConfig:
    def __init__(self, config_item: dict):
        self.id = config_item.get("id")
        self.describe = config_item.get("describe")
        self.login_user = config_item.get("login_user")
        self.login_password = config_item.get("login_password")
        self.ssh_private_key= config_item.get("ssh_private_key")
        self.host_status = {host: "IDLE" for host in config_item.get("hosts")}


@app.post("ask_nodes")
def ask_nodes_fun():
    """
        请求空闲节点
    :return: dict is_idle: bool; idle_hosts: list
    """

    pass


def init_hosts_status():
    """
    读取machine_pool.json，扫描机器使用情况，初始化资源池
    """
    with open(MACHINE_POOL_FILE) as f:
        machine_pool = json.load(f)
        for item in machine_pool.items():
            # todo 扫描资源池
            pass


class StandaloneApplication(BaseApplication):
    def __init__(self, app, options=None):
        self.options = options or {}
        self.application = app
        super().__init__()

    def load_config(self):
        for key, value in self.options.items():
            if key in self.cfg.settings and value is not None:
                self.cfg.set(key.lower(), value)

    def load(self):
        return self.application


if __name__ == "__main__":
    options = {
        'bind': f'0.0.0.0:4253',
        'workers': 8,
        'timeout': 600,
        'worker_class': 'uvicorn.workers.UvicornWorker',
    }

    standalone_app = StandaloneApplication(app, options)
    standalone_app.run()